#include "GPUMemoryManager.h"
#include <cuda.h>
#include <cuda_runtime.h>
#include "cuda_errorcheck.h"
#include "cuda_defs.h"
#include "cuda_rrt_types.h"
#include "PlanningParams.h"

#define DEBUG_GPU_MEMORY_MANAGER

using namespace std;

GpuMemoryManager* GpuMemoryManager::instance = NULL;

GpuMemoryManager::GpuMemoryManager()
{

}

GpuMemoryManager::~GpuMemoryManager()
{
	destroy();
}

GpuMemoryManager* GpuMemoryManager::getInstance()
{
	if (instance == NULL)
		instance = new GpuMemoryManager();
	return instance;
}

void GpuMemoryManager::init()
{
	// problem specific
	addNewChunk(sizeof(SerializedStates) * NUM_THREADS);
	addNewChunk(sizeof(float) * NUM_THREADS);
	addNewChunk(sizeof(int) * NUM_THREADS);
	addNewChunk(sizeof(float) * NUM_THREADS * MAX_SAMPLES);
	addNewChunk(sizeof(int) * NUM_THREADS);
	addNewChunk(sizeof(float) * NUM_THREADS);
	addNewChunk(sizeof(float) * NUM_THREADS);
	addNewChunk(sizeof(SerializedStates) * NUM_THREADS);
	addNewChunk(sizeof(float) * NUM_THREADS * MAX_SAMPLES);
	addNewChunk(sizeof(float)* NUM_THREADS * (MAX_NUM_NEAREST_NEIGHBOR + 512));
	addNewChunk(sizeof(int)* NUM_THREADS * (MAX_NUM_NEAREST_NEIGHBOR + 512));
	addNewChunk(sizeof(float)* NUM_THREADS * (MAX_NUM_NEAREST_NEIGHBOR + 512));
	addNewChunk(sizeof(float)* NUM_THREADS * (MAX_NUM_NEAREST_NEIGHBOR + 512));
}

void GpuMemoryManager::destroy()
{
	std::multiset<MemoryChunk*, MemoryChunkSizeComp>::iterator it = memorySet.begin();
	for (; it != memorySet.end(); ++it)
	{
		MemoryChunk* chunk = *it;
		GPUFREE(chunk->p);
		delete chunk;
	}
	memorySet.clear();
}

GpuMemoryManager::MemoryChunk* GpuMemoryManager::getChunk(unsigned int size)
{
	MemoryChunk temp;
	temp.size = size;

	std::multiset<MemoryChunk*, MemoryChunkSizeComp>:: iterator it = memorySet.lower_bound(&temp);
	for (; it != memorySet.end(); ++it)
	{
		MemoryChunk* chunk = *it;
		if (chunk->free)
		{
			chunk->free = false;
			return chunk;
		}
	}

	MemoryChunk* chunk = addNewChunk(size);
	chunk->free = false;
	return chunk;
}

void GpuMemoryManager::releaseChunk(MemoryChunk* memory)
{
	memory->free = true;
}

GpuMemoryManager::MemoryChunk* GpuMemoryManager::addNewChunk(unsigned int size)
{
	MemoryChunk* chunk = new MemoryChunk;
	chunk->free = true;
	chunk->size = size;
	GPUMALLOC(&chunk->p, size);
	memorySet.insert(chunk);

#ifdef DEBUG_GPU_MEMORY_MANAGER
	printf("[GPUMM] New memory allocaled : %d\n", size);
#endif

	return chunk;
}